{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# python 端 AnnotateUsedMemory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/media/pc/data/lxw/ai/tvm-book/doc/read/relay\n"
     ]
    }
   ],
   "source": [
    "%cd ..\n",
    "import testing"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```cpp\n",
    "size_t DivRoundUp(size_t size, size_t word_size) { return (size + word_size - 1) / word_size; }\n",
    "\n",
    "size_t GetMemorySizeBytes(const Array<PrimExpr>& shape, const DataType& dtype) {\n",
    "  size_t size = 1;\n",
    "  for (IndexExpr dim : shape) {\n",
    "    const int64_t* pval = tir::as_const_int(dim);\n",
    "    ICHECK(pval != nullptr) << \"Cannot allocate memory symbolic tensor shape \" << shape;\n",
    "    ICHECK_GE(*pval, 0) << \"Cannot allocate memory for tensor with negative shape\" << *pval;\n",
    "    size *= static_cast<size_t>(pval[0]);\n",
    "  }\n",
    "  size *= DivRoundUp(dtype.bits() * dtype.lanes(), 8);\n",
    "  return size;\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tvm.runtime import DataType\n",
    "\n",
    "def div_round_up(size, word_size): # ceil(size/word_size) DivRoundUp 向上取整\n",
    "    \"\"\"计算一个给定大小（size）的元素需要多少个指定大小的单词（word_size）来存储。\n",
    "    具体来说，它通过将原始大小除以元素大小并向上取整来实现这一点。\n",
    "\n",
    "    args:\n",
    "        - `size`：表示原始大小，即需要存储的元素的大小。\n",
    "        - `word_size`：表示每个单词的大小，即每个单词可以存储的元素数量。\n",
    "\n",
    "    Returns:\n",
    "        - 返回值为整数，表示需要的单词数量。\n",
    "    \"\"\"\n",
    "    val = (size + word_size - 1) // word_size\n",
    "    return val\n",
    "\n",
    "def get_memory_size_bytes(shape, dtype): # GetMemorySizeBytes\n",
    "    if isinstance(dtype, str):\n",
    "        dtype = DataType(dtype)\n",
    "    assert isinstance(dtype, DataType)\n",
    "    size = 1\n",
    "    for pval in shape:\n",
    "        size *= pval\n",
    "    # dtype.bits 表示数据类型的位数\n",
    "    # dtype.lanes 表示数据类型的通道数\n",
    "    size *= div_round_up(dtype.bits * dtype.lanes, 8)\n",
    "    return size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(1, 16)"
      ]
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dtype = DataType(\"int8\")\n",
    "div_round_up(dtype.bits * dtype.lanes, 8), get_memory_size_bytes((1, 2, 2, 4), dtype)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tvm\n",
    "from tvm import relay\n",
    "from tvm.relay.expr_functor import ExprVisitor\n",
    "\n",
    "\n",
    "def AnnotateUsedMemory():\n",
    "    return relay.transform._ffi_api.AnnotateUsedMemory()\n",
    "\n",
    "class CheckUsedMemoryAnnotation(ExprVisitor):\n",
    "    \"\"\"\n",
    "    Check that the annotations on each function in the graph match\n",
    "    what is expected.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, expected_annotations, expected_io_annotation):\n",
    "        self.expected_annotations = expected_annotations\n",
    "        self.expected_io_annotation = expected_io_annotation\n",
    "        super().__init__()\n",
    "\n",
    "    def visit_function(self, fn):\n",
    "        if \"Primitive\" in fn.attrs:\n",
    "            assert (\n",
    "                \"used_memory\" in fn.attrs\n",
    "            ), \"Primitive function does not have used_memory annotation.\"\n",
    "\n",
    "            assert len(self.expected_annotations) > 0, \"Not all expected annotations were compared\"\n",
    "\n",
    "            expected_mem = self.expected_annotations.pop(0)\n",
    "            actual_mem = [int(x) for x in fn.attrs[\"used_memory\"]]\n",
    "            assert expected_mem == actual_mem, (\n",
    "                f\"Expected used memory annotation {expected_mem} \"\n",
    "                f\"did not match actual annotation {actual_mem}\"\n",
    "            )\n",
    "        super().visit_function(fn)\n",
    "\n",
    "    def __call__(self, fn):\n",
    "        assert (\n",
    "            fn.attrs[\"io_used_memory\"] == self.expected_io_annotation\n",
    "        ), \"Expected IO annotation did not match.\"\n",
    "        self.visit(fn.body)\n",
    "\n",
    "\n",
    "def _check_used_memory_annotations(mod, expected_annotations, expected_io_annotation):\n",
    "    mod = relay.transform.InferType()(mod)\n",
    "    mod = relay.transform.ToANormalForm()(mod)\n",
    "    mod = relay.transform.InferType()(mod)\n",
    "    mod = AnnotateUsedMemory()(mod)\n",
    "    CheckUsedMemoryAnnotation(expected_annotations, expected_io_annotation)(mod[\"main\"])\n",
    "    relay.transform.ToGraphNormalForm()(mod).show()\n",
    "\n",
    "def _create_primitive_function(expr):\n",
    "    func = relay.Function(relay.analysis.free_vars(expr), expr)\n",
    "    func = func.with_attr(\"Primitive\", 1)\n",
    "    return func"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 注解简单原语函数内存"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, io_used_memory<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">32</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), int8] {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, Primitive<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, used_memory<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">32</span>]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), int8] {\n",
       "    nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>max_pool2d(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x, pool_size<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">1</span>], padding<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>]) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), int8]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def get_inner_func():\n",
    "    x = relay.var(\"x\", shape=(1, 2, 2, 4), dtype=\"int8\")\n",
    "    x = relay.nn.max_pool2d(x)\n",
    "    x = _create_primitive_function(x)\n",
    "    return x\n",
    "\n",
    "ifm = relay.var(\"input\", shape=(1, 2, 2, 4), dtype=\"int8\")\n",
    "call = relay.Call(get_inner_func(), [ifm])\n",
    "mod = tvm.IRModule.from_expr(call)\n",
    "\n",
    "expected_annotations = [\n",
    "    [2 * (1 * 2 * 2 * 4)],\n",
    "]\n",
    "expected_io_annotation = 2 * (1 * 2 * 2 * 4)\n",
    "_check_used_memory_annotations(mod, expected_annotations, expected_io_annotation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">2</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">2</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, io_used_memory<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">178</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">2</span>), int8] {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>x2: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">2</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">2</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, Primitive<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, used_memory<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">226</span>]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">2</span>), int8] {\n",
       "    nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>max_pool2d(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x2, pool_size<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>], padding<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], layout<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;NHWC&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">2</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">8</span>, <span style=\"color: #008000\">2</span>), int8]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">2</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">2</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>x1: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">2</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">2</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, Primitive<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, used_memory<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">170</span>]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">6</span>, <span style=\"color: #008000\">6</span>, <span style=\"color: #008000\">2</span>), int8] {\n",
       "    nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>max_pool2d(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x1, pool_size<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>], padding<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], layout<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;NHWC&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">6</span>, <span style=\"color: #008000\">6</span>, <span style=\"color: #008000\">2</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">7</span>, <span style=\"color: #008000\">2</span>), int8]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">6</span>, <span style=\"color: #008000\">6</span>, <span style=\"color: #008000\">2</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">6</span>, <span style=\"color: #008000\">6</span>, <span style=\"color: #008000\">2</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">6</span>, <span style=\"color: #008000\">6</span>, <span style=\"color: #008000\">2</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">6</span>, <span style=\"color: #008000\">6</span>, <span style=\"color: #008000\">2</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, Primitive<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, used_memory<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">122</span>]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">2</span>), int8] {\n",
       "    nn<span style=\"color: #AA22FF; font-weight: bold\">.</span>max_pool2d(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x, pool_size<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>], padding<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>, <span style=\"color: #008000\">0</span>], layout<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;NHWC&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">2</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">6</span>, <span style=\"color: #008000\">6</span>, <span style=\"color: #008000\">2</span>), int8]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">2</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">2</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def get_inner_func(ifm_shape):\n",
    "    x = relay.var(\"x\", shape=ifm_shape, dtype=\"int8\")\n",
    "    x = relay.nn.max_pool2d(x, pool_size=(2, 2), layout=\"NHWC\")\n",
    "    x = _create_primitive_function(x)\n",
    "    return x\n",
    "\n",
    "ifm = relay.var(\"input\", shape=(1, 8, 8, 2), dtype=\"int8\")\n",
    "x = get_inner_func((1, 8, 8, 2))\n",
    "x = relay.Call(x, [ifm])\n",
    "y = get_inner_func((1, 7, 7, 2))\n",
    "y = relay.Call(y, [x])\n",
    "z = get_inner_func((1, 6, 6, 2))\n",
    "z = relay.Call(z, [y])\n",
    "mod = tvm.IRModule.from_expr(z)\n",
    "\n",
    "expected_annotations = [\n",
    "    [(1 * 8 * 8 * 2) + (1 * 7 * 7 * 2)],\n",
    "    [(1 * 7 * 7 * 2) + (1 * 6 * 6 * 2)],\n",
    "    [(1 * 6 * 6 * 2) + (1 * 5 * 5 * 2)],\n",
    "]\n",
    "expected_io_annotation = (1 * 8 * 8 * 2) + (1 * 5 * 5 * 2)\n",
    "_check_used_memory_annotations(mod, expected_annotations, expected_io_annotation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 注解混合数据类型的原语函数内存"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>), int16] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>), int16] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, io_used_memory<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">48</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>), uint32] {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>), int16] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>), int16] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, Primitive<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, used_memory<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">48</span>]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>), uint32] {\n",
       "    cast(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x, dtype<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;uint32&quot;</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>), uint32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>), int16]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>), uint32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">2</span>), uint32] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def get_inner_func():\n",
    "    x = relay.var(\"x\", shape=(1, 2, 2, 2), dtype=\"int16\")\n",
    "    x = relay.cast(x, dtype=\"uint32\")\n",
    "    x = _create_primitive_function(x)\n",
    "    return x\n",
    "\n",
    "ifm = relay.var(\"input\", shape=(1, 2, 2, 2), dtype=\"int16\")\n",
    "x = get_inner_func()\n",
    "x = relay.Call(x, [ifm])\n",
    "mod = tvm.IRModule.from_expr(x)\n",
    "\n",
    "expected_annotations = [\n",
    "    [(1 * 2 * 2 * 2) * 2 + (1 * 2 * 2 * 2) * 4],\n",
    "]\n",
    "expected_io_annotation = (1 * 2 * 2 * 2) * 2 + (1 * 2 * 2 * 2) * 4\n",
    "_check_used_memory_annotations(mod, expected_annotations, expected_io_annotation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 注解 `concatenate` 内存"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "当将两个函数的结果连接成一个结果时，第二个函数还将具有第一个函数的活跃结果，因此将使用较大的“已使用内存”值进行标注。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #AA22FF\">@main</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">6</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">6</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, io_used_memory<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">360</span>) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">30</span>), int8] {\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>x: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">6</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">6</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, Primitive<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, used_memory<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">240</span>]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">30</span>), int8] {\n",
       "    reshape(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x, newshape<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">30</span>]) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">30</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">6</span>), int8]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">30</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> fn (<span style=\"color: #AA22FF; font-weight: bold\">%</span>x1: Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">6</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">6</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>, Primitive<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #008000\">1</span>, used_memory<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">360</span>]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">30</span>), int8] {\n",
       "    reshape(<span style=\"color: #AA22FF; font-weight: bold\">%</span>x1, newshape<span style=\"color: #AA22FF; font-weight: bold\">=</span>[<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">30</span>]) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">30</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "  } <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>fn (Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">6</span>), int8]) <span style=\"color: #AA22FF; font-weight: bold\">-&gt;</span> Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">30</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">0</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">30</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">1</span>(<span style=\"color: #AA22FF; font-weight: bold\">%</span>input) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">30</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span> <span style=\"color: #AA22FF; font-weight: bold\">=</span> (<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">2</span>, <span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">3</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>(Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">30</span>), int8], Tensor[(<span style=\"color: #008000\">1</span>, <span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">30</span>), int8]) <span style=\"color: #AA22FF; font-weight: bold\">*/</span>;\n",
       "  concatenate(<span style=\"color: #AA22FF; font-weight: bold\">%</span><span style=\"color: #008000\">4</span>) <span style=\"color: #AA22FF; font-weight: bold\">/*</span> ty<span style=\"color: #AA22FF; font-weight: bold\">=</span>Tensor[(<span style=\"color: #008000\">2</span>, <span style=\"color: #008000\">4</span>, <span style=\"color: #008000\">30</span>), int8] <span style=\"color: #AA22FF; font-weight: bold\">*/</span>\n",
       "}\n",
       "</pre></div>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "def get_inner_func():\n",
    "    x = relay.var(\"x\", shape=(1, 4, 5, 6), dtype=\"int8\")\n",
    "    x = relay.reshape(x, newshape=(1, 4, 30))\n",
    "    x = _create_primitive_function(x)\n",
    "    return x\n",
    "\n",
    "ifm = relay.var(\"input\", shape=(1, 4, 5, 6), dtype=\"int8\")\n",
    "x = relay.Call(get_inner_func(), [ifm])\n",
    "y = relay.Call(get_inner_func(), [ifm])\n",
    "z = relay.concatenate([x, y], axis=0)\n",
    "mod = tvm.IRModule.from_expr(z)\n",
    "\n",
    "expected_annotations = [\n",
    "    [(1 * 4 * 5 * 6) + (1 * 4 * 30)],\n",
    "    # the output tensor from the previous function is also alive\n",
    "    [(1 * 4 * 5 * 6) + (1 * 4 * 30) + (1 * 4 * 30)],\n",
    "]\n",
    "expected_io_annotation = (1 * 4 * 5 * 6) + (1 * 4 * 60)\n",
    "_check_used_memory_annotations(mod, expected_annotations, expected_io_annotation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 注解其他数据"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "<>:2: SyntaxWarning: invalid escape sequence '\\ '\n",
      "<>:171: SyntaxWarning: invalid escape sequence '\\ '\n",
      "<>:2: SyntaxWarning: invalid escape sequence '\\ '\n",
      "<>:171: SyntaxWarning: invalid escape sequence '\\ '\n",
      "/tmp/ipykernel_448750/3780910861.py:2: SyntaxWarning: invalid escape sequence '\\ '\n",
      "  \"\"\"\n",
      "/tmp/ipykernel_448750/3780910861.py:171: SyntaxWarning: invalid escape sequence '\\ '\n",
      "  \"\"\"\n"
     ]
    }
   ],
   "source": [
    "def test_many_different_parallel_calls():\n",
    "    \"\"\"\n",
    "    Test a graph that calls many different functions in parallel.\n",
    "\n",
    "                    input\n",
    "            /         |         \\\n",
    "    prim_func_1  prim_func_2  prim_func_3\n",
    "           \\         |         /\n",
    "                 prim_func_4\n",
    "    \"\"\"\n",
    "\n",
    "    def get_inner_func_1():\n",
    "        x = relay.var(\"x\", shape=(1, 4, 5, 6), dtype=\"int8\")\n",
    "        x = relay.tanh(x)\n",
    "        x = _create_primitive_function(x)\n",
    "        return x\n",
    "\n",
    "    def get_inner_func_2():\n",
    "        x = relay.var(\"x\", shape=(1, 4, 5, 6), dtype=\"int8\")\n",
    "        x = relay.nn.max_pool2d(x, pool_size=(1, 1), layout=\"NHWC\")\n",
    "        x = _create_primitive_function(x)\n",
    "        return x\n",
    "\n",
    "    def get_inner_func_3():\n",
    "        x = relay.var(\"x\", shape=(1, 4, 5, 6), dtype=\"int8\")\n",
    "        x = relay.abs(x)\n",
    "        x = relay.nn.relu(x)\n",
    "        x = relay.exp(x)\n",
    "        x = _create_primitive_function(x)\n",
    "        return x\n",
    "\n",
    "    def get_inner_func_4():\n",
    "        x = relay.var(\"x\", shape=(1, 4, 5, 6), dtype=\"int8\")\n",
    "        y = relay.var(\"y\", shape=(1, 4, 5, 6), dtype=\"int8\")\n",
    "        z = relay.var(\"z\", shape=(1, 4, 5, 6), dtype=\"int8\")\n",
    "        out = relay.concatenate([x, y, z], axis=3)\n",
    "        out = _create_primitive_function(out)\n",
    "        return out\n",
    "\n",
    "    ifm = relay.var(\"input\", shape=(1, 4, 5, 6), dtype=\"int8\")\n",
    "    x = relay.Call(get_inner_func_1(), [ifm])\n",
    "    y = relay.Call(get_inner_func_2(), [ifm])\n",
    "    z = relay.Call(get_inner_func_3(), [ifm])\n",
    "    a = relay.Call(get_inner_func_4(), [x, y, z])\n",
    "    mod = tvm.IRModule.from_expr(a)\n",
    "\n",
    "    expected_annotations = [\n",
    "        [(1 * 4 * 5 * 6) + (1 * 4 * 5 * 6)],\n",
    "        # output from prim_func_1 is also still alive\n",
    "        [(1 * 4 * 5 * 6) + (1 * 4 * 5 * 6) + (1 * 4 * 5 * 6)],\n",
    "        # outputs from prim_func_1 and prim_func_2 are also still alive\n",
    "        [(1 * 4 * 5 * 6) + (1 * 4 * 5 * 6) + (1 * 4 * 5 * 6) + (1 * 4 * 5 * 6)],\n",
    "        [(1 * 4 * 5 * 6) + (1 * 4 * 5 * 6) + (1 * 4 * 5 * 6) + (1 * 4 * 5 * 18)],\n",
    "    ]\n",
    "    expected_io_annotation = (1 * 4 * 5 * 6) + (1 * 4 * 5 * 18)\n",
    "    _check_used_memory_annotations(mod, expected_annotations, expected_io_annotation)\n",
    "\n",
    "\n",
    "def test_nested_branches():\n",
    "    \"\"\"\n",
    "    Tests a graph with branches that also branch.\n",
    "\n",
    "             input\n",
    "            /     \\\n",
    "          /        \\\n",
    "    prim_func_1  prim_func_2\n",
    "                   /     \\\n",
    "                  /       \\\n",
    "            prim_func_3   prim_func_4\n",
    "    \"\"\"\n",
    "\n",
    "    def get_generic_inner_func():\n",
    "        x = relay.var(\"x\", shape=(1, 2, 2, 4), dtype=\"int8\")\n",
    "        x = relay.nn.relu(x)\n",
    "        return _create_primitive_function(x)\n",
    "\n",
    "    ifm = relay.var(\"input\", shape=(1, 2, 2, 4), dtype=\"int8\")\n",
    "    a = relay.Call(get_generic_inner_func(), [ifm])\n",
    "    b = relay.Call(get_generic_inner_func(), [ifm])\n",
    "    c = relay.Call(get_generic_inner_func(), [b])\n",
    "    d = relay.Call(get_generic_inner_func(), [b])\n",
    "    out = relay.concatenate([a, c, d], axis=3)\n",
    "    mod = tvm.IRModule.from_expr(out)\n",
    "\n",
    "    expected_annotations = [\n",
    "        [(1 * 2 * 2 * 4) + (1 * 2 * 2 * 4)],\n",
    "        # output from prim_func_1 is also still alive\n",
    "        [(1 * 2 * 2 * 4) + (1 * 2 * 2 * 4) + (1 * 2 * 2 * 4)],\n",
    "        # output from prim_func_1 is also still alive\n",
    "        [(1 * 2 * 2 * 4) + (1 * 2 * 2 * 4) + (1 * 2 * 2 * 4)],\n",
    "        # outputs from prim_func_1 and prim_func_3 are also still alive\n",
    "        [(1 * 2 * 2 * 4) + (1 * 2 * 2 * 4) + (1 * 2 * 2 * 4) + (1 * 2 * 2 * 4)],\n",
    "    ]\n",
    "    expected_io_annotation = (1 * 2 * 2 * 4) + (1 * 2 * 2 * 12)\n",
    "    _check_used_memory_annotations(mod, expected_annotations, expected_io_annotation)\n",
    "\n",
    "\n",
    "def test_composite_inner_function():\n",
    "    \"\"\"\n",
    "    Tests the typical BYOC use case where a primitive function\n",
    "    contains a composite function.\n",
    "    \"\"\"\n",
    "\n",
    "    def get_inner_func():\n",
    "        x = relay.var(\"x\", shape=(1, 2, 2, 4), dtype=\"int8\")\n",
    "        x = relay.nn.max_pool2d(x, pool_size=(2, 2), layout=\"NHWC\")\n",
    "        x = relay.Function(relay.analysis.free_vars(x), x)\n",
    "        x = x.with_attr(\"Composite\", \"my_composite_func\")\n",
    "\n",
    "        y = relay.var(\"y\", shape=(1, 2, 2, 4), dtype=\"int8\")\n",
    "        z = relay.Call(x, [y])\n",
    "        return _create_primitive_function(z)\n",
    "\n",
    "    ifm = relay.var(\"input\", shape=(1, 2, 2, 4), dtype=\"int8\")\n",
    "    x = relay.Call(get_inner_func(), [ifm])\n",
    "    mod = tvm.IRModule.from_expr(x)\n",
    "\n",
    "    expected_annotations = [\n",
    "        [(1 * 2 * 2 * 4) + (1 * 1 * 1 * 4)],\n",
    "    ]\n",
    "    expected_io_annotation = (1 * 2 * 2 * 4) + (1 * 1 * 1 * 4)\n",
    "    _check_used_memory_annotations(mod, expected_annotations, expected_io_annotation)\n",
    "\n",
    "\n",
    "def test_multiple_calls_to_same_function():\n",
    "    \"\"\"\n",
    "    Tests the case when there are multiple calls to the same function.\n",
    "    \"\"\"\n",
    "\n",
    "    def get_inner_func():\n",
    "        x = relay.var(\"x\", shape=(1, 2, 2, 4), dtype=\"int8\")\n",
    "        x = relay.nn.max_pool2d(x)\n",
    "        x = _create_primitive_function(x)\n",
    "        return x\n",
    "\n",
    "    inner_func = get_inner_func()\n",
    "    ifm = relay.var(\"input\", shape=(1, 2, 2, 4), dtype=\"int8\")\n",
    "    call1 = relay.Call(inner_func, [ifm])\n",
    "    call2 = relay.Call(inner_func, [call1])\n",
    "    mod = tvm.IRModule.from_expr(call2)\n",
    "\n",
    "    expected_annotations = [[2 * (1 * 2 * 2 * 4), 2 * (1 * 2 * 2 * 4)]]\n",
    "    expected_io_annotation = 2 * (1 * 2 * 2 * 4)\n",
    "    _check_used_memory_annotations(mod, expected_annotations, expected_io_annotation)\n",
    "\n",
    "\n",
    "def test_parallel_calls_to_same_function():\n",
    "    \"\"\"\n",
    "    Test parallel calls to the same function.\n",
    "    \"\"\"\n",
    "\n",
    "    def get_inner_func():\n",
    "        x = relay.var(\"x\", shape=(1, 2, 2, 4), dtype=\"int8\")\n",
    "        x = relay.nn.max_pool2d(x)\n",
    "        x = _create_primitive_function(x)\n",
    "        return x\n",
    "\n",
    "    inner_func = get_inner_func()\n",
    "    ifm = relay.var(\"input\", shape=(1, 2, 2, 4), dtype=\"int8\")\n",
    "    call1 = relay.Call(inner_func, [ifm])\n",
    "    call2 = relay.Call(inner_func, [ifm])\n",
    "    concat = relay.concatenate([call1, call2], axis=0)\n",
    "    mod = tvm.IRModule.from_expr(concat)\n",
    "\n",
    "    expected_annotations = [[2 * (1 * 2 * 2 * 4), 3 * (1 * 2 * 2 * 4)]]\n",
    "    expected_io_annotation = 3 * (1 * 2 * 2 * 4)\n",
    "    _check_used_memory_annotations(mod, expected_annotations, expected_io_annotation)\n",
    "\n",
    "\n",
    "def test_parallel_calls_with_non_ifm_input():\n",
    "    \"\"\"\n",
    "    Test a graph that calls many different functions in parallel where\n",
    "    the input is not the input to the function.\n",
    "\n",
    "                    y = f(x)\n",
    "            /         |         \\\n",
    "       z0 = g0(y)    ...      zi = gi(y)\n",
    "           \\         |         /\n",
    "                  concat\n",
    "    \"\"\"\n",
    "\n",
    "    def get_inner_func_1():\n",
    "        x = relay.var(\"x\", shape=(1, 4, 5, 6), dtype=\"int8\")\n",
    "        x = relay.tanh(x)\n",
    "        x = _create_primitive_function(x)\n",
    "        return x\n",
    "\n",
    "    def get_inner_func_2():\n",
    "        x = relay.var(\"x\", shape=(1, 4, 5, 6), dtype=\"int8\")\n",
    "        x = relay.nn.max_pool2d(x, pool_size=(2, 2))\n",
    "        x = _create_primitive_function(x)\n",
    "        return x\n",
    "\n",
    "    ifm = relay.var(\"input\", shape=(1, 4, 5, 6), dtype=\"int8\")\n",
    "    y = relay.Call(get_inner_func_1(), [ifm])\n",
    "    g = get_inner_func_2()\n",
    "\n",
    "    no_calls = 20\n",
    "    z = [relay.Call(g, [y]) for _ in range(0, no_calls)]\n",
    "    out = relay.concatenate(z, axis=3)\n",
    "    mod = tvm.IRModule.from_expr(out)\n",
    "\n",
    "    expected_annotations = [\n",
    "        [(1 * 4 * 5 * 6) + (1 * 4 * 5 * 6)],\n",
    "        [(1 * 4 * 5 * 6) + (1 * 4 * 4 * 5) * i for i in range(1, no_calls + 1)],\n",
    "    ]\n",
    "    expected_io_annotation = (1 * 4 * 5 * 6) + (1 * 4 * 4 * (5 * no_calls))\n",
    "    _check_used_memory_annotations(mod, expected_annotations, expected_io_annotation)\n",
    "\n",
    "\n",
    "def test_dynamic_io_tensor_not_supported():\n",
    "    \"\"\"\n",
    "    Test to check dynamic IO tensor error.\n",
    "    \"\"\"\n",
    "\n",
    "    def get_inner_func():\n",
    "        x = relay.var(\"x\", shape=(1, 2, 2, 4), dtype=\"int8\")\n",
    "        x = relay.nn.max_pool2d(x)\n",
    "        x = _create_primitive_function(x)\n",
    "        return x\n",
    "\n",
    "    ifm = relay.var(\"input\", shape=(1, 2, 2, relay.Any()), dtype=\"int8\")\n",
    "    call = relay.Call(get_inner_func(), [ifm])\n",
    "    mod = tvm.IRModule.from_expr(call)\n",
    "\n",
    "    err_rgx = r\"AnnotateUsedMemory does not support dynamic shapes\"\n",
    "    with pytest.raises(tvm.TVMError, match=err_rgx):\n",
    "        _check_used_memory_annotations(mod, [], [])\n",
    "\n",
    "\n",
    "def test_dynamic_callsite_tensor_not_supported():\n",
    "    \"\"\"\n",
    "    Test to check dynamic callsite tensor error.\n",
    "    \"\"\"\n",
    "\n",
    "    def get_inner_func():\n",
    "        x = relay.var(\"x\", shape=(relay.Any(), 2, 2, 4), dtype=\"int8\")\n",
    "        x = relay.nn.max_pool2d(x)\n",
    "        x = _create_primitive_function(x)\n",
    "        return x\n",
    "\n",
    "    ifm = relay.var(\"input\", shape=(1, 2, 2, 4), dtype=\"int8\")\n",
    "    call = relay.Call(get_inner_func(), [ifm])\n",
    "    mod = tvm.IRModule.from_expr(call)\n",
    "\n",
    "    err_rgx = r\"AnnotateUsedMemory does not support dynamic shapes\"\n",
    "    with pytest.raises(tvm.TVMError, match=err_rgx):\n",
    "        _check_used_memory_annotations(mod, [], [])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "xinet",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
